import os
os.environ["CUDA_VISIBLE_DEVICES"] = "0"
import sys
sys.path.append(".")
import time
from pathlib import Path
from loguru import logger
from datetime import datetime

from hyvideo.utils.file_utils import save_videos_grid
from hyvideo.config import parse_args
from hyvideo.inference import HunyuanVideoSampler


def main(start_idx=0, end_idx=152):
    args = parse_args()
    print(args)
    models_root_path = Path(args.model_base)
    if not models_root_path.exists():
        raise ValueError(f"`models_root` not exists: {models_root_path}")
    
    # Create save folder to save the samples
    save_path = args.save_path if args.save_path_suffix=="" else f'{args.save_path}_{args.save_path_suffix}'
    if not os.path.exists(save_path):
        os.makedirs(save_path, exist_ok=True)

    # Load models
    hunyuan_video_sampler = HunyuanVideoSampler.from_pretrained(models_root_path, args=args)
    
    # Get the updated args
    args = hunyuan_video_sampler.args

    # Start sampling
    # TODO: batch inference check

    # 多次顺序生成跑批
    # data_root = "/share/home/wuqingyao_danglingwei/datas/TACO_Data_20250314/full_20k_plus"
    data_root = "/WORK/PUBLIC/liuyebin_work/lingweidang/datas/TACO_Data_20250314/full_20k_plus"
    save_dir = r"hunyuani2v13B_samples_416_624"
    os.makedirs(save_dir, exist_ok=True)
    img_column_path = "test_small_151_images_416x624.txt"
    prompt_column_path = "test_small_151_VLMEnhanced_prompts.txt" # "test_small_151_VLMEnhanced_prompts.txt" "test_small_151_prompts.txt"
    img_path_list = []
    prompts_list = []
    with open(os.path.join(data_root, img_column_path), "r", encoding="utf-8") as f:
        all_lines = f.readlines()

    for line in all_lines:
        line = line.strip()
        img_path_list.append(line)

    with open(os.path.join(data_root, prompt_column_path), "r", encoding="utf-8") as f:
        all_lines = f.readlines()

    for line in all_lines:
        line = line.strip()
        prompts_list.append(line)

    assert len(img_path_list) == len(prompts_list), f"img_path_list and prompts_list should have the same length."

    for idx, (ref_image_path, prompt) in enumerate(zip(img_path_list, prompts_list)):
        if not start_idx <= idx < end_idx:
            continue

        ref_image_path = os.path.join(data_root, ref_image_path) # first_frame_416_624_00000
        data_sort_idx = int(os.path.splitext(os.path.basename(ref_image_path))[0].split("_")[-1])
        print(f"idx: {idx}, ref_image_path: {ref_image_path}")

        outputs = hunyuan_video_sampler.predict(
            prompt=prompt, 
            i2v_image_path=ref_image_path,
            height=args.video_size[0],
            width=args.video_size[1],
            video_length=args.video_length,
            seed=args.seed,
            negative_prompt=args.neg_prompt,
            infer_steps=args.infer_steps,
            guidance_scale=args.cfg_scale,
            num_videos_per_prompt=args.num_videos,
            flow_shift=args.flow_shift,
            batch_size=args.batch_size,
            embedded_guidance_scale=args.embedded_cfg_scale,
            i2v_mode=args.i2v_mode,
            i2v_resolution=args.i2v_resolution,
            i2v_condition_type=args.i2v_condition_type,
            i2v_stability=args.i2v_stability,
            ulysses_degree=args.ulysses_degree,
            ring_degree=args.ring_degree,
        )
        samples = outputs['samples']
        print(f"生成结果：{samples[0].shape}") # [3, 49, 416, 624]
        # Save samples
        if 'LOCAL_RANK' not in os.environ or int(os.environ['LOCAL_RANK']) == 0:
            for i, sample in enumerate(samples):
                save_path = os.path.join(save_dir, f"pred_color_videos_{sample.shape[-2]}x{sample.shape[-1]}_{data_sort_idx:05d}.mp4") 
                sample = samples[i].unsqueeze(0)
                # time_flag = datetime.fromtimestamp(time.time()).strftime("%Y-%m-%d-%H:%M:%S")
                # cur_save_path = f"{save_path}/{time_flag}_seed{outputs['seeds'][i]}_{outputs['prompts'][i][:100].replace('/','')}.mp4"
                save_videos_grid(sample, save_path, fps=8)
                logger.info(f'Sample save to: {save_path}')

if __name__ == "__main__":
    main(start_idx=0, end_idx=152)
